#include "AnalyticsBlocker.h"
#include "Debug.h"
#include "Assertion.h"
#include "../Managers/Hook.h"
#include <algorithm>

#ifdef _WIN64
#pragma comment(lib, "ws2_32.lib")
#else
#pragma comment(lib, "wsock32.lib")
#endif

bool AnalyticsBlocker::ShouldDAB = false;
const char* AnalyticsBlocker::BlockedList[] = { new char[8] { 0x68, 0x61, 0x74, 0x65, 0x62, 0x69, 0x6e, 0x00 }, new char[14] { 0x6f, 0x63, 0x75, 0x6c, 0x75, 0x73, 0x63, 0x64, 0x6e, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[9] { 0x70, 0x61, 0x73, 0x74, 0x65, 0x62, 0x69, 0x6e, 0x00 }, new char[24] { 0x61, 0x6e, 0x61, 0x6c, 0x79, 0x74, 0x69, 0x63, 0x73, 0x2e, 0x6e, 0x69, 0x6e, 0x6a, 0x61, 0x6b, 0x69, 0x77, 0x69, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[22] { 0x66, 0x61, 0x63, 0x65, 0x62, 0x6f, 0x6f, 0x6b, 0x2d, 0x68, 0x61, 0x72, 0x64, 0x77, 0x61, 0x72, 0x65, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[18] { 0x67, 0x61, 0x6d, 0x65, 0x61, 0x6e, 0x61, 0x6c, 0x79, 0x74, 0x69, 0x63, 0x73, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[10] { 0x66, 0x62, 0x63, 0x64, 0x6e, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[12] { 0x69, 0x63, 0x65, 0x62, 0x75, 0x72, 0x6e, 0x2e, 0x78, 0x79, 0x7a, 0x00 }, new char[10] { 0x66, 0x62, 0x73, 0x62, 0x78, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[7] { 0x66, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[21] { 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2d, 0x61, 0x6e, 0x61, 0x6c, 0x79, 0x74, 0x69, 0x63, 0x73, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[25] { 0x70, 0x75, 0x62, 0x6c, 0x69, 0x63, 0x2e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[30] { 0x70, 0x65, 0x72, 0x66, 0x2d, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[8] { 0x6d, 0x65, 0x61, 0x70, 0x2e, 0x67, 0x67, 0x00 }, new char[8] { 0x64, 0x72, 0x6f, 0x70, 0x62, 0x6f, 0x78, 0x00 }, new char[29] { 0x63, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x2e, 0x75, 0x63, 0x61, 0x2e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[14] { 0x61, 0x6d, 0x70, 0x6c, 0x69, 0x74, 0x75, 0x64, 0x65, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[10] { 0x66, 0x62, 0x63, 0x64, 0x6e, 0x2e, 0x6e, 0x65, 0x74, 0x00 }, new char[26] { 0x65, 0x63, 0x6f, 0x6d, 0x6d, 0x65, 0x72, 0x63, 0x65, 0x2e, 0x69, 0x61, 0x70, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[42] { 0x64, 0x61, 0x74, 0x61, 0x2d, 0x6f, 0x70, 0x74, 0x6f, 0x75, 0x74, 0x2d, 0x73, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x2e, 0x75, 0x63, 0x61, 0x2e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[13] { 0x76, 0x72, 0x6d, 0x6f, 0x64, 0x73, 0x2e, 0x73, 0x70, 0x61, 0x63, 0x65, 0x00 }, new char[6] { 0x66, 0x62, 0x2e, 0x6d, 0x65, 0x00 }, new char[21] { 0x70, 0x69, 0x78, 0x65, 0x6c, 0x73, 0x74, 0x72, 0x69, 0x6b, 0x65, 0x33, 0x64, 0x61, 0x77, 0x73, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[17] { 0x69, 0x6e, 0x67, 0x65, 0x73, 0x74, 0x2e, 0x73, 0x65, 0x6e, 0x74, 0x72, 0x79, 0x2e, 0x69, 0x6f, 0x00 }, new char[8] { 0x64, 0x69, 0x73, 0x63, 0x6f, 0x72, 0x64, 0x00 }, new char[21] { 0x67, 0x6c, 0x75, 0x65, 0x68, 0x65, 0x6e, 0x64, 0x65, 0x72, 0x2d, 0x61, 0x6c, 0x75, 0x68, 0x75, 0x74, 0x2e, 0x64, 0x65, 0x00 }, new char[9] { 0x68, 0x61, 0x73, 0x74, 0x65, 0x62, 0x69, 0x6e, 0x00 }, new char[16] { 0x63, 0x72, 0x61, 0x73, 0x68, 0x6c, 0x79, 0x74, 0x69, 0x63, 0x73, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[26] { 0x61, 0x70, 0x69, 0x2e, 0x75, 0x63, 0x61, 0x2e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[13] { 0x66, 0x61, 0x63, 0x65, 0x62, 0x6f, 0x6f, 0x6b, 0x2e, 0x6e, 0x65, 0x74, 0x00 }, new char[9] { 0x73, 0x6b, 0x69, 0x64, 0x2d, 0x68, 0x75, 0x62, 0x00 }, new char[22] { 0x63, 0x64, 0x70, 0x2e, 0x63, 0x6c, 0x6f, 0x75, 0x64, 0x2e, 0x75, 0x6e, 0x69, 0x74, 0x79, 0x33, 0x64, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[13] { 0x72, 0x69, 0x70, 0x70, 0x65, 0x72, 0x2e, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x00 }, new char[18] { 0x61, 0x7a, 0x75, 0x72, 0x65, 0x77, 0x65, 0x62, 0x73, 0x69, 0x74, 0x65, 0x73, 0x2e, 0x6e, 0x65, 0x74, 0x00 }, new char[11] { 0x69, 0x74, 0x65, 0x6f, 0x2e, 0x73, 0x70, 0x61, 0x63, 0x65, 0x00 }, new char[16] { 0x73, 0x6f, 0x66, 0x74, 0x6c, 0x69, 0x67, 0x68, 0x74, 0x2e, 0x61, 0x74, 0x2e, 0x75, 0x61, 0x00 }, new char[13] { 0x66, 0x61, 0x63, 0x65, 0x62, 0x6f, 0x6f, 0x6b, 0x2e, 0x63, 0x6f, 0x6d, 0x00 } }; 
std::list<std::string> AnalyticsBlocker::DABList = { new char[8] { 0x6e, 0x74, 0x70, 0x2e, 0x6f, 0x72, 0x67, 0x00 }, new char[13] { 0x62, 0x6f, 0x6e, 0x65, 0x74, 0x6f, 0x6d, 0x65, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[11] { 0x73, 0x61, 0x6d, 0x62, 0x6f, 0x79, 0x2e, 0x64, 0x65, 0x76, 0x00 }, new char[11] { 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[14] { 0x72, 0x75, 0x62, 0x79, 0x2d, 0x63, 0x6f, 0x72, 0x65, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[16] { 0x6d, 0x65, 0x6c, 0x6f, 0x6e, 0x6c, 0x6f, 0x61, 0x64, 0x65, 0x72, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[22] { 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x75, 0x73, 0x65, 0x72, 0x63, 0x6f, 0x6e, 0x74, 0x65, 0x6e, 0x74, 0x2e, 0x63, 0x6f, 0x6d, 0x00 }, new char[20] { 0x74, 0x68, 0x65, 0x74, 0x72, 0x75, 0x65, 0x79, 0x6f, 0x73, 0x68, 0x69, 0x66, 0x61, 0x6e, 0x2e, 0x63, 0x6f, 0x6d, 0x00 } };

bool AnalyticsBlocker::Initialize()
{
	Debug::Msg("Initializing Analytics Blocker...");
	return wsock32::Initialize()
#ifdef _WIN64
		&& ws2_32::Initialize()
#endif
;}

void AnalyticsBlocker::Hook()
{
	wsock32::Hooks::Initialize();
#ifdef _WIN64
	ws2_32::Hooks::Initialize();
#endif
}

bool AnalyticsBlocker::CheckHostNameOrIP(const char* host_name_or_ip)
{
	std::string host_name_or_ip_str = host_name_or_ip;
	std::transform(host_name_or_ip_str.begin(), host_name_or_ip_str.end(), host_name_or_ip_str.begin(), [](unsigned char c) { return std::tolower(c); });
	if (host_name_or_ip_str._Equal("localhost") || host_name_or_ip_str._Equal("127.0.0.1"))
		return false;
	bool should_block = ShouldBlock(host_name_or_ip);
	if (Debug::Enabled || ShouldDAB)
	{
		if (should_block)
			Debug::DirectWrite(("Host Name or IP Blocked: " + host_name_or_ip_str).c_str());
		else if (!HasDabbed(host_name_or_ip))
		{
			Debug::DirectWrite(("Unique Host Name or IP Found: " + host_name_or_ip_str).c_str());
			DABList.push_back(host_name_or_ip_str);
		}
	}
	return should_block;
}

bool AnalyticsBlocker::ShouldBlock(const char* host_name_or_ip)
{
	bool found = false;
	for (int i = 0; i < (sizeof(BlockedList) / sizeof(BlockedList[0])); i++)
	{
		if (strstr(host_name_or_ip, BlockedList[i]) != NULL)
		{
			found = true;
			break;
		}
	}
	return found;
}

bool AnalyticsBlocker::HasDabbed(const char* host_name_or_ip)
{
	bool found = false;
	for (auto entry = DABList.begin(); entry != DABList.end(); ++entry)
	{
		if ((*entry)._Equal(host_name_or_ip))
		{
			found = true;
			break;
		}
	}
	return found;
}

#pragma region wsock32
HMODULE AnalyticsBlocker::wsock32::Module = NULL;
AnalyticsBlocker::wsock32::Exports::gethostbyname_t AnalyticsBlocker::wsock32::Exports::Gethostbyname = NULL;

bool AnalyticsBlocker::wsock32::Initialize()
{
	Debug::Msg("Initializing wsock32...");
	Module = LoadLibraryA("wsock32.dll");
	if (Module != NULL)
	{
		if (!Exports::Initialize())
			return false;
	}
	else
	{
		// Display Warning
	}
	return true;
}

bool AnalyticsBlocker::wsock32::Exports::Initialize()
{
	Debug::Msg("Initializing wsock32 Exports...");

	Gethostbyname = (gethostbyname_t)Assertion::GetExport(Module, "gethostbyname");

	return Assertion::ShouldContinue;
}

void AnalyticsBlocker::wsock32::Hooks::Initialize()
{
	if (Module == NULL)
		return;
	Debug::Msg("Attaching Hooks to wsock32...");

	Debug::Msg("gethostbyname");
	Hook::Attach(&(LPVOID&)Exports::Gethostbyname, Gethostbyname);
}

hostent* AnalyticsBlocker::wsock32::Hooks::Gethostbyname(const char* name)
{
	try
	{
		if (name != NULL && CheckHostNameOrIP(name))
			name = "localhost";
		return Exports::Gethostbyname(name);
	}
	catch(...){}
	WSASetLastError(WSATRY_AGAIN);
	return NULL;
}
#pragma endregion

#ifdef _WIN64
#pragma region ws2_32
HMODULE AnalyticsBlocker::ws2_32::Module = NULL;
AnalyticsBlocker::ws2_32::Exports::getaddrinfo_t AnalyticsBlocker::ws2_32::Exports::Getaddrinfo = NULL;

bool AnalyticsBlocker::ws2_32::Initialize()
{
	Debug::Msg("Initializing ws2_32...");
	Module = LoadLibraryA("ws2_32");
	if (Module != NULL)
	{
		if (!Exports::Initialize())
			return false;
	}
	else
	{
		// Display Warning
	}
	return true;
}

bool AnalyticsBlocker::ws2_32::Exports::Initialize()
{
	Debug::Msg("Initializing ws2_32 Exports...");

	Getaddrinfo = (getaddrinfo_t)Assertion::GetExport(Module, "getaddrinfo");

	return Assertion::ShouldContinue;
}

void AnalyticsBlocker::ws2_32::Hooks::Initialize()
{
	if (Module == NULL)
		return;
	Debug::Msg("Attaching Hooks to ws2_32...");

	Debug::Msg("getaddrinfo");
	Hook::Attach(&(LPVOID&)Exports::Getaddrinfo, Getaddrinfo);
}

DWORD AnalyticsBlocker::ws2_32::Hooks::Getaddrinfo(PCSTR pNodeName, PCSTR pServiceName, const void* pHints, void* ppResult)
{
	try
	{
		if (pNodeName != NULL && CheckHostNameOrIP(pNodeName))
			pNodeName = "localhost";
		return Exports::Getaddrinfo(pNodeName, pServiceName, pHints, ppResult);
	}
	catch (...){}
	WSASetLastError(WSATRY_AGAIN);
	return WSATRY_AGAIN;
}
#pragma endregion
#endif